package com.bruce.ai.alibaba.agentcard.discovery.circuitbreaker;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Duration;
import java.time.Instant;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.function.Supplier;

/**
 * 熔断器实现
 * 防止级联故障，当错误率超过阈值时自动熔断
 */
public class CircuitBreaker {
    
    private static final Logger logger = LoggerFactory.getLogger(CircuitBreaker.class);
    
    private final String name;
    private final CircuitBreakerConfig config;
    private final Executor executor;
    
    // 状态管理
    private final AtomicReference<State> state = new AtomicReference<>(State.CLOSED);
    private final AtomicLong lastFailureTime = new AtomicLong(0);
    private final AtomicInteger failureCount = new AtomicInteger(0);
    private final AtomicInteger successCount = new AtomicInteger(0);
    private final AtomicLong lastStateChangeTime = new AtomicLong(System.currentTimeMillis());
    
    // 统计信息
    private final AtomicLong totalRequests = new AtomicLong(0);
    private final AtomicLong totalFailures = new AtomicLong(0);
    private final AtomicLong totalSuccesses = new AtomicLong(0);
    
    public CircuitBreaker(String name, CircuitBreakerConfig config) {
        this(name, config, ForkJoinPool.commonPool());
    }
    
    public CircuitBreaker(String name, CircuitBreakerConfig config, Executor executor) {
        this.name = name;
        this.config = config;
        this.executor = executor;
    }
    
    /**
     * 熔断器状态
     */
    public enum State {
        CLOSED,     // 关闭状态，正常处理请求
        OPEN,       // 开启状态，拒绝所有请求
        HALF_OPEN   // 半开状态，允许少量请求测试服务是否恢复
    }
    
    /**
     * 执行操作
     * @param operation 要执行的操作
     * @param <T> 返回类型
     * @return 操作结果
     * @throws Exception 如果操作失败或熔断器开启
     */
    public <T> T execute(Callable<T> operation) throws Exception {
        return execute(operation, null);
    }
    
    /**
     * 执行操作，支持回退策略
     * @param operation 要执行的操作
     * @param fallback 回退策略
     * @param <T> 返回类型
     * @return 操作结果
     * @throws Exception 如果操作失败且没有回退策略
     */
    public <T> T execute(Callable<T> operation, Function<Exception, T> fallback) throws Exception {
        totalRequests.incrementAndGet();
        
        State currentState = getCurrentState();
        
        switch (currentState) {
            case OPEN:
                logger.debug("Circuit breaker [{}] is OPEN, rejecting request", name);
                CircuitBreakerOpenException openException = 
                    new CircuitBreakerOpenException("Circuit breaker is OPEN");
                
                if (fallback != null) {
                    return fallback.apply(openException);
                } else {
                    throw openException;
                }
                
            case HALF_OPEN:
                logger.debug("Circuit breaker [{}] is HALF_OPEN, allowing test request", name);
                return executeInHalfOpenState(operation, fallback);
                
            case CLOSED:
            default:
                return executeInClosedState(operation, fallback);
        }
    }
    
    /**
     * 异步执行操作
     * @param operation 要执行的操作
     * @param <T> 返回类型
     * @return CompletableFuture包装的结果
     */
    public <T> CompletableFuture<T> executeAsync(Supplier<T> operation) {
        return executeAsync(operation, null);
    }
    
    /**
     * 异步执行操作，支持回退策略
     * @param operation 要执行的操作
     * @param fallback 回退策略
     * @param <T> 返回类型
     * @return CompletableFuture包装的结果
     */
    public <T> CompletableFuture<T> executeAsync(Supplier<T> operation, Function<Exception, T> fallback) {
        return CompletableFuture.supplyAsync(() -> {
            try {
                return execute(() -> operation.get(), fallback);
            } catch (Exception e) {
                if (e instanceof RuntimeException) {
                    throw (RuntimeException) e;
                } else {
                    throw new RuntimeException(e);
                }
            }
        }, executor);
    }
    
    /**
     * 在关闭状态下执行操作
     */
    private <T> T executeInClosedState(Callable<T> operation, Function<Exception, T> fallback) throws Exception {
        try {
            T result = operation.call();
            onSuccess();
            return result;
        } catch (Exception e) {
            onFailure();
            
            if (fallback != null) {
                return fallback.apply(e);
            } else {
                throw e;
            }
        }
    }
    
    /**
     * 在半开状态下执行操作
     */
    private <T> T executeInHalfOpenState(Callable<T> operation, Function<Exception, T> fallback) throws Exception {
        try {
            T result = operation.call();
            onSuccessInHalfOpen();
            return result;
        } catch (Exception e) {
            onFailureInHalfOpen();
            
            if (fallback != null) {
                return fallback.apply(e);
            } else {
                throw e;
            }
        }
    }
    
    /**
     * 处理成功情况
     */
    private void onSuccess() {
        successCount.incrementAndGet();
        totalSuccesses.incrementAndGet();
        
        // 重置失败计数
        if (failureCount.get() > 0) {
            failureCount.set(0);
        }
    }
    
    /**
     * 处理失败情况
     */
    private void onFailure() {
        int failures = failureCount.incrementAndGet();
        totalFailures.incrementAndGet();
        lastFailureTime.set(System.currentTimeMillis());
        
        // 检查是否需要开启熔断器
        if (shouldTripCircuitBreaker(failures)) {
            tripCircuitBreaker();
        }
    }
    
    /**
     * 在半开状态下处理成功
     */
    private void onSuccessInHalfOpen() {
        logger.info("Circuit breaker [{}] test request succeeded, closing circuit breaker", name);
        closeCircuitBreaker();
        onSuccess();
    }
    
    /**
     * 在半开状态下处理失败
     */
    private void onFailureInHalfOpen() {
        logger.warn("Circuit breaker [{}] test request failed, reopening circuit breaker", name);
        tripCircuitBreaker();
        onFailure();
    }
    
    /**
     * 检查是否应该开启熔断器
     */
    private boolean shouldTripCircuitBreaker(int failures) {
        return failures >= config.getFailureThreshold();
    }
    
    /**
     * 开启熔断器
     */
    private void tripCircuitBreaker() {
        if (state.compareAndSet(State.CLOSED, State.OPEN) || 
            state.compareAndSet(State.HALF_OPEN, State.OPEN)) {
            
            lastStateChangeTime.set(System.currentTimeMillis());
            logger.warn("Circuit breaker [{}] tripped to OPEN state after {} failures", 
                       name, failureCount.get());
        }
    }
    
    /**
     * 关闭熔断器
     */
    private void closeCircuitBreaker() {
        if (state.compareAndSet(State.HALF_OPEN, State.CLOSED)) {
            lastStateChangeTime.set(System.currentTimeMillis());
            failureCount.set(0);
            successCount.set(0);
            logger.info("Circuit breaker [{}] closed", name);
        }
    }
    
    /**
     * 获取当前状态
     */
    private State getCurrentState() {
        State currentState = state.get();
        
        if (currentState == State.OPEN) {
            // 检查是否可以转换到半开状态
            long timeSinceLastFailure = System.currentTimeMillis() - lastStateChangeTime.get();
            if (timeSinceLastFailure >= config.getWaitDurationInOpenState().toMillis()) {
                if (state.compareAndSet(State.OPEN, State.HALF_OPEN)) {
                    lastStateChangeTime.set(System.currentTimeMillis());
                    logger.info("Circuit breaker [{}] transitioned from OPEN to HALF_OPEN", name);
                    return State.HALF_OPEN;
                }
            }
        }
        
        return currentState;
    }
    
    /**
     * 获取熔断器状态
     */
    public State getState() {
        return getCurrentState();
    }
    
    /**
     * 获取熔断器名称
     */
    public String getName() {
        return name;
    }
    
    /**
     * 获取配置
     */
    public CircuitBreakerConfig getConfig() {
        return config;
    }
    
    /**
     * 获取统计信息
     */
    public CircuitBreakerMetrics getMetrics() {
        return new CircuitBreakerMetrics(
            totalRequests.get(),
            totalSuccesses.get(),
            totalFailures.get(),
            failureCount.get(),
            successCount.get(),
            getState(),
            Duration.ofMillis(System.currentTimeMillis() - lastStateChangeTime.get())
        );
    }
    
    /**
     * 手动重置熔断器
     */
    public void reset() {
        state.set(State.CLOSED);
        failureCount.set(0);
        successCount.set(0);
        lastFailureTime.set(0);
        lastStateChangeTime.set(System.currentTimeMillis());
        logger.info("Circuit breaker [{}] manually reset", name);
    }
    
    /**
     * 手动开启熔断器
     */
    public void forceOpen() {
        state.set(State.OPEN);
        lastStateChangeTime.set(System.currentTimeMillis());
        logger.warn("Circuit breaker [{}] manually forced to OPEN", name);
    }
    
    /**
     * 熔断器开启异常
     */
    public static class CircuitBreakerOpenException extends Exception {
        public CircuitBreakerOpenException(String message) {
            super(message);
        }
    }
}